# -*- coding: utf-8 -*-
import pandas as pd
import json
import os
import logging
import re

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger("convert_restaurant_data")

# 文件路径
current_dir = os.path.dirname(os.path.abspath(__file__))
csv_path = os.path.join(current_dir, '大众点评.csv')
knowledge_path = os.path.join(current_dir, 'knowledge.json')

def load_existing_knowledge():
    """加载现有的知识库文件"""
    if os.path.exists(knowledge_path):
        try:
            with open(knowledge_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            logger.error("知识库文件格式错误，将创建新文件")
            return []
        except Exception as e:
            logger.error(f"读取知识库文件时出错: {e}")
            return []
    else:
        logger.info("知识库文件不存在，将创建新文件")
        return []

def clean_text(text):
    """清理文本，去除多余的引号、空格和换行符"""
    if isinstance(text, str):
        # 去除引号和额外空格
        text = text.strip('"\'').strip()
        # 将连续的空格替换为单个空格
        text = re.sub(r'\s+', ' ', text)
        return text
    return text

def clean_dishes(dishes):
    """清理招牌菜文本，返回列表"""
    if not isinstance(dishes, str):
        return []
    
    # 去除引号
    dishes = dishes.strip('"\'')
    
    # 分割菜品（可能用分号、逗号等分隔）
    items = re.split(r'[;,，、]', dishes)
    
    # 清理每个菜品名称并过滤空项
    return [item.strip() for item in items if item.strip()]

def extract_score(score_text):
    """从评分文本中提取数值"""
    if isinstance(score_text, str):
        # 尝试提取数字
        match = re.search(r'(\d+(\.\d+)?)', score_text)
        if match:
            return float(match.group(1))
    return None

def extract_price(price_text):
    """从价格文本中提取数值"""
    if isinstance(price_text, str):
        # 尝试提取数字
        match = re.search(r'(\d+(\.\d+)?)', price_text)
        if match:
            return int(float(match.group(1)))
    return None

def convert_to_knowledge_entries():
    """将餐厅数据转换为知识条目"""
    try:
        # 读取CSV文件
        logger.info(f"正在读取CSV文件: {csv_path}")
        df = pd.read_csv(csv_path, encoding='utf-8')
        logger.info(f"CSV文件读取成功，共 {len(df)} 条记录")
        
        # 清理数据
        df = df.applymap(clean_text)
        logger.info("数据清理完成")
        
        # 加载现有知识条目
        knowledge_data = load_existing_knowledge()
        initial_count = len(knowledge_data)
        logger.info(f"已加载现有知识库，包含 {initial_count} 条知识条目")
        
        # 获取最大ID
        max_id = 0
        if knowledge_data:
            max_id = max(item['id'] for item in knowledge_data)
        
        # 创建知识条目
        new_entries = []
        city_restaurants = {}
        cuisine_restaurants = {}
        
        # 1. 为每个餐厅创建详细知识条目
        for _, row in df.iterrows():
            # 跳过空行或格式不正确的行
            if pd.isna(row['门店名称']) or not row['门店名称'].strip():
                continue
                
            max_id += 1
            
            # 提取基本信息
            city = row['城市']
            name = row['门店名称']
            address = row['地址']
            cuisine = row['品类']
            hours = row['营业时间']
            overall_score = extract_score(row['门店整体评分'])
            taste_score = extract_score(row['口味评分'])
            price = extract_price(row['人均消费'])
            signature_dishes = clean_dishes(row['招牌菜'])
            
            # 添加到城市和菜系分类
            if city not in city_restaurants:
                city_restaurants[city] = []
            city_restaurants[city].append(name)
            
            if cuisine not in cuisine_restaurants:
                cuisine_restaurants[cuisine] = []
            cuisine_restaurants[cuisine].append(name)
            
            # 创建详细描述
            description = f"{name}是位于{city}的一家{cuisine}餐厅，地址：{address}"
            
            if hours:
                description += f"，营业时间：{hours}"
            
            description += "。"
            
            if overall_score:
                description += f"该餐厅的整体评分为{overall_score}分"
                
                if taste_score:
                    description += f"，口味评分为{taste_score}分"
                
                description += "。"
            
            if price:
                description += f"人均消费约{price}元。"
            
            # 添加招牌菜信息
            if signature_dishes:
                # 如果招牌菜太多，只取前10个
                if len(signature_dishes) > 10:
                    top_dishes = signature_dishes[:10]
                    description += f"招牌菜包括：{', '.join(top_dishes)}等。"
                else:
                    description += f"招牌菜包括：{', '.join(signature_dishes)}。"
            
            # 添加关键词部分
            keywords = f"关键词：{city} {name} {cuisine} 餐厅 美食"
            if signature_dishes and len(signature_dishes) > 0:
                keywords += f" {signature_dishes[0]}"
            if price:
                if price < 50:
                    keywords += " 经济实惠"
                elif price > 200:
                    keywords += " 高档餐厅"
            
            # 组合完整知识条目
            restaurant_entry = {
                "id": max_id,
                "text": f"{description}\n\n{keywords}"
            }
            
            new_entries.append(restaurant_entry)
        
        # 2. 创建城市美食汇总
        for city, restaurants in city_restaurants.items():
            max_id += 1
            
            city_text = f"{city}美食指南\n\n{city}共有{len(restaurants)}家餐厅数据，包括："
            
            # 限制显示餐厅数量
            if len(restaurants) > 15:
                city_text += "、".join(restaurants[:15]) + "等。"
            else:
                city_text += "、".join(restaurants) + "。"
                
            city_text += f"\n\n关键词：{city} 美食 餐厅 指南 推荐"
            
            city_entry = {
                "id": max_id,
                "text": city_text
            }
            
            new_entries.append(city_entry)
        
        # 3. 创建菜系美食汇总
        for cuisine, restaurants in cuisine_restaurants.items():
            max_id += 1
            
            cuisine_text = f"{cuisine}美食指南\n\n数据中共有{len(restaurants)}家{cuisine}餐厅，包括："
            
            # 限制显示餐厅数量
            if len(restaurants) > 15:
                cuisine_text += "、".join(restaurants[:15]) + "等。"
            else:
                cuisine_text += "、".join(restaurants) + "。"
                
            cuisine_text += f"\n\n关键词：{cuisine} 美食 餐厅 指南 推荐"
            
            cuisine_entry = {
                "id": max_id,
                "text": cuisine_text
            }
            
            new_entries.append(cuisine_entry)
            
        # 添加到知识库
        knowledge_data.extend(new_entries)
        logger.info(f"已创建 {len(new_entries)} 条新知识条目")
        
        # 保存更新后的知识库
        with open(knowledge_path, 'w', encoding='utf-8') as f:
            json.dump(knowledge_data, f, ensure_ascii=False, indent=2)
        
        logger.info(f"知识库已更新，当前共有 {len(knowledge_data)} 条知识条目")
        logger.info(f"新增 {len(knowledge_data) - initial_count} 条知识条目")
        
        # 打印一条示例条目
        if new_entries:
            logger.info(f"示例条目: {new_entries[0]['text'][:200]}...")
        
        return True
    
    except Exception as e:
        logger.error(f"转换过程中出错: {e}")
        import traceback
        logger.error(traceback.format_exc())
        return False

if __name__ == "__main__":
    logger.info("开始转换餐厅数据到知识库")
    success = convert_to_knowledge_entries()
    if success:
        logger.info("转换完成")
    else:
        logger.error("转换失败") 